# coding: utf-8

import os
import tensorflow as tf
import tensorflow.contrib.keras as kr

from rnn_model import TRNNConfig, TextRNN
from data.cnews_loader import read_category, read_vocab

# try:
#     bool(type(unicode))
# except NameError:
#     unicode = str
#
# base_dir = 'data/cnews'
# vocab_dir = os.path.join(base_dir, 'cnews.vocab.txt')
#
# save_dir = 'checkpoints/textrcnn'
# save_path = os.path.join(save_dir, 'best_validation')  # 最佳验证结果保存路径


save_path = './checkpoints/textrnn/best_validation'
# import os
# os.environ['CUDA_VISIBLE_DEVICES'] = "1"

class RnnModel:
    def __init__(self):
        self.config = TRNNConfig()

        # self.categories, self.cat_to_id = read_category()
        # self.words, self.word_to_id = read_vocab(vocab_dir)
        self.config.vocab_size = 6000
        self.model = TextRNN(self.config)

        self.session = tf.Session()
        self.session.run(tf.compat.v1.global_variables_initializer())
        saver = tf.compat.v1.train.Saver()
        saver.restore(sess=self.session, save_path=save_path)  # 读取保存的模型

    def predict(self, message):
        # 支持不论在python2还是python3下训练的模型都可以在2或者3的环境下运行
        # content = unicode(message)
        # data = [self.word_to_id[x] for x in content if x in self.word_to_id]

        feed_dict = {
            self.model.input_x: kr.preprocessing.sequence.pad_sequences([message], self.config.seq_length),
            self.model.keep_prob: 1.0
        }

        y_pred_cls = self.session.run(self.model.y_pred_cls, feed_dict=feed_dict)
        return y_pred_cls[0]



rnn_model = RnnModel()
# test_demo = ['三星ST550以全新的拍摄方式超越了以往任何一款数码相机',
# '热火vs骑士前瞻：皇帝回乡二番战 东部次席唾手可得新浪体育讯北京时间3月30日7:00']
# for i in test_demo:
#     print(cnn_model.predict(i))
# test_demo = '三星ST550以全新的拍摄方式超越了以往任何一款数码相机'
# test_demo = [97, 72, 123, 15, 1, 38, 101]
# test_demo = [50,33,51,20,0.9,33,63]#21
test_demo = [61,42,72,19,0.9,44,54]

print(rnn_model.predict(test_demo))
